# -*- coding:utf-8 -*-
from bert4keras.tokenizers import Tokenizer
from chat_bot.models.intent.bert_model import build_model

# 类别
class_nums = 21

# 路径
best_model_filepath = 'model/best_model.weights'
# 预训练模型路径
checkpoint_path = 'model/bert_model.ckpt'
config_path = 'model/bert_config.json'
dict_path = 'model/vocab.txt'


class IntentModel(object):
    def __init__(self):
        super(IntentModel, self).__init__()
        self.dict_path = dict_path
        self.config_path = config_path
        self.checkpoint_path = checkpoint_path
        self.label_list = [line.strip() for line in open(work_path + 'data/label', 'r', encoding='utf8')]
        self.tokenizer = Tokenizer(self.dict_path)
        self.model = build_model(self.config_path, self.checkpoint_path, class_nums)
        self.model.load_weights(best_model_filepath)

    def predict(self, text):
        # 转化问句输入
        token_ids, segment_ids = self.tokenizer.encode(text, maxlen=60)
        # 得到每个分类的预测概率
        probability_list = self.model.predict([[token_ids], [segment_ids]])
        # 将每个分类与其概率对应起来
        predict = {label: probability for label, probability in zip(self.label_list, probability_list[0])}
        # 以字典的值进行排序
        predict = sorted(predict.items(), key=lambda kv: kv[1], reverse=True)
        # 拿概率最高的分类
        label, confidence = predict[0]
        return {"label": label, "confidence": float(confidence)}

# if __name__ == '__main__':
#     intentModel = IntentModel()
#     while True:
#         q = input(":")
#         print(intentModel.predict(q))
